'''
@Time    : 2022/3/24 22:04
@Author  : Fu Junyu
@Site    : www.fujunyu.cn
@File    : train_RGCN.py
@Software: PyCharm
'''
import pickle
import torch
from torch import nn, optim
from datalodar import datalodar, trainData, userFeatTest, moviesFeatTest
from model import RGCNModel
import cmath
from matplotlib import pyplot as plt


def MAE(a, b):
    c = b - a
    return torch.abs(c).float().mean(0)

def train(epochs, f1=False):
    # 创建图形
    test = datalodar('ML')
    graph = test.forward()

    # 获取用户特征
    user_max_dict = {
        'uid': 943,  # 用户ID的最大值
        'gender': 2,	 # 0/1 表示男女
        'age': 74,     # 7个年龄段
        'job': 21,     # 21个工作分类
        'zip': 99836
    }
    movie_max_dict = {
        'mid': 1682,  # 电影ID的最大值
        'mtype': 19,  # 电影类型的最大值
    }
    convParams = {
        'kernel_sizes': [2, 3, 4, 5]
    }

    userfeat = userFeatTest(dataset='ML')
    feat_user = userfeat.getFeat

    # 获取电影特征
    moviesfeat = moviesFeatTest(dataset='ML')
    feat_movies = moviesfeat.getFeat

    feat_dic = {'user': feat_user, 'movies': feat_movies}

    # 创建含RGCN层的模型
    model = RGCNModel(5, 788, 20, user_max_dict, movie_max_dict, 5, 12)

# --------------------------------训练参数设置---------------------------------------
    LEARNING_RATE = 0.001
    WEIGHT_DACAY = 5e-4
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    if f1:
        loss = torch.nn.CrossEntropyLoss().to(DEVICE)
    else:
        loss = nn.MSELoss().to(DEVICE)
        loss_test = nn.MSELoss().to(DEVICE)
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DACAY)

    train_result = []
    test_result = []
    test_MAE = []


    for epoch in range(epochs):
        trainAccu = []
        testAccu = []
        testMae = []

        # 获取训练数据集
        traindata = trainData(dataset='ML')
        traindata = traindata.getData

        # 获取测试数据集
        testdata = trainData(dataset='ML_test')
        testdata = testdata.getData

        print(f'epoch: {epoch}')
        for data in traindata:

            label_pred = model(graph, feat_dic, data)

            l = loss(label_pred, data['labels'])
            optimizer.zero_grad()
            l.backward()
            optimizer.step()

            trainAccu.append(cmath.sqrt(l.item()).real)

        for t in testdata:
            l_test = loss_test(model(graph, feat_dic, t), t['labels'])
            mae = MAE(model(graph, feat_dic, t), t['labels'])
            testAccu.append(cmath.sqrt(l_test.item()).real)
            testMae.append(mae.detach().numpy())


        test_result.append(min(testAccu))
        train_result.append(min(trainAccu))
        test_MAE.append(min(testMae))

        print(f'train RMSE:{min(trainAccu)}, test RMSE:{min(testAccu)}, test MAE:{min(testMae)}')

    print('train best RMSE:', min(train_result), 'test best RMSE:', min(test_result), 'test best MAE:', min(testMae))
    torch.save(model.state_dict(), 'Params/model_params.pkl')

    f = open('Params/RGCN_RMSE.pkl', 'wb')
    pickle.dump(test_result, f)
    f.close()

    f = open('Params/RGCN_MAE.pkl', 'wb')
    pickle.dump(test_MAE, f)
    f.close()

    x = range(1, 101)
    y1 = train_result
    y2 = test_result
    Y3 = test_MAE
    plt.plot(x, y1, color='magenta', linestyle='--', label='train_RMSE')
    plt.plot(x, y2, color='blue', label='test_RMSE')
    plt.plot(x, Y3, color='green', label='MAE')

    plt.xlabel('epochs')
    plt.ylabel('RMSE/MAE')

    plt.legend(loc=0)
    xmin, xmax, ymin, ymax = 0, 101, 0.6, 1.31
    plt.axis([xmin, xmax, ymin, ymax])

    plt.savefig("Params/test.svg", dpi=300, format="svg")
    plt.show()


if __name__ == '__main__':
    train(100)






